# coding=utf-8
import tushare as ts
import json
from datetime import datetime, timedelta
import dateutil
from peewee import *
import os
from models import Index, Flag, IndexDaily, db
from sqlalchemy import create_engine
from models import utils
import talib as ta
import pandas as pd
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
import matplotlib.dates as mdate

engine = create_engine('sqlite:///data.sqlite')

ts.set_token('b46ae864716dbb76999f35738dcd531a8eab6809740045dd78b8ceba')
pro = ts.pro_api()

markets = ['MSCI', 'CSI', 'SSE', 'SZSE', 'CICC', 'SW', 'CNI', 'OTH']

# for m in markets:
#     print 'm', m
#     df = pro.index_basic(market=m)
#     df.to_sql('index', con=engine, if_exists='append', index=False)

# ts.get_today_all()

# print ts.get_h_data('39
# 9106', index=True, start='2015-01-01', end='2015-03-16')


indexes = [
    '000300.SH',  # 300
    '000905.SH',  # 500
    '000016.SH',  # 50
    '399006.SZ',  # 创业板
    '000990.SH',  # 全指消费
    '000807.SH',  # 食品饮料
    '000993.SH',  # 全指信息
    'h11136.CSI', # 中国互联
    '399934.SZ',  # 中证金融
]

today = datetime.now()

### index_daily

max_dates = IndexDaily.select(
    IndexDaily.ts_code,
    fn.max(IndexDaily.trade_date).alias('trade_date'),
    ).group_by(IndexDaily.ts_code)

max_dates_dict = {}
for row in max_dates:
    max_dates_dict[row.ts_code] = row.trade_date

print max_dates_dict

for idx in indexes:
    max_date = utils.date2str(
        utils.str2date(max_dates_dict[idx]) + timedelta(days=1)
    )
    print 'max_date', max_date
    # df = pro.index_daily(ts_code=idx, start_date=max_date)
    # df.to_sql('index_daily', con=engine, if_exists='append', index=False)

    df = pd.read_sql_query('select * from main.index_daily where ts_code="%s"' % idx, con=engine)

    rsi = ta.RSI(df.close.values)
    rsi6 = ta.RSI(df.close.values,6)
    rsi30 = ta.RSI(df.close.values,30)
    cci, cci6, cci30 = None,None,None
    try:
        cci = ta.CCI(df.high.values, df.low.values, df.close.values)
        cci6 = ta.CCI(df.high.values, df.low.values, df.close.values, timeperiod=6)
        cci30 = ta.CCI(df.high.values, df.low.values, df.close.values, timeperiod=30)
    except Exception as e:
        print idx, e
    # new_df = df.loc[df.trade_date > '20180810', ['trade_date', 'rsi']]
    new_df = df.loc[df.trade_date > utils.date2str(today)]
    print 'new_df', len(new_df)
    for values in (rsi, rsi6, rsi30, cci, cci6, cci30):
        if values is not None:
            for i, v in enumerate(values):
                values[i] = round(v, 3)
    old_len = len(df) - len(new_df)
    for i, close in db.batch_commit(enumerate(new_df.close.values), 500):
        i = i + old_len
        if cci is not None:
            update = IndexDaily.update(
                rsi=rsi[i], rsi6=rsi6[i],
                rsi30=rsi30[i], cci=cci[i],
                cci6=cci6[i], cci30=cci30[i]
            )
        else:
            update = IndexDaily.update(
                rsi=rsi[i], rsi6=rsi6[i],
                rsi30=rsi30[i],
            )
        trade_date = df.trade_date.values[i]
        update = update.where(IndexDaily.ts_code == idx and IndexDaily.trade_date == trade_date)
        update.execute()

    # break
    rsi_df = pd.DataFrame(data={'rsi': rsi, 'rsi6': rsi6, 'rsi30': rsi30}, index=map(utils.str2date, df.trade_date.values))
    # print rsi_df

    buy_point = rsi_df.loc[(rsi_df.rsi6 < 15) & (rsi_df.rsi30 < 40)]
    print 'buy_point', buy_point

    sell_point = rsi_df.loc[(rsi_df.rsi6 > 80) & (rsi_df.rsi30 > 70)]
    print 'sell_point', sell_point

    p = rsi_df.plot()
    p.xaxis.set_major_formatter(mdate.DateFormatter('%Y-%m-%d'))
    plt.show()


### index_weight
# for idx in indexes:
#     idx_info = Index.select().where(Index.ts_code == idx)[0]
#     print 'ts_code', idx_info.ts_code
#     start_date = utils.str2date(idx_info.base_date)

#     cur_date = today
#     while cur_date >= start_date:
#         iw_df = pro.index_weight(index_code=idx_info.ts_code, trade_date=utils.date2str(cur_date))
#         iw_df.to_sql('index_weight', con=engine, if_exists='append', index=False)
#         print iw_df
#         print ('query ts_code:%s date: %s' %  (idx_info.ts_code, utils.date2str(cur_date)))
#         cur_date -= timedelta(days=1)


# df = pro.index_weight(index_code='000016.SH', trade_date='20180928')

# print df

# pro.daily_basic(ts_code='', trade_date='20080726')
# print pro.query('daily_basic', ts_code='801002', trade_date='20180721', fields='pe,pb,ts_code,close')

# print pro.index_basic(market='SW')
# print pro.index_basic(market='CSI')
# df = pro.index_weight(index_code='399300.SZ', trade_date='20180903')

# print df.to_json()
